package cool.doudou.doudada.cipher.algorithm.util;

import cool.doudou.doudada.cipher.algorithm.enums.Algorithm;

import javax.crypto.Cipher;
import java.nio.charset.StandardCharsets;
import java.security.KeyFactory;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.SecureRandom;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;
import java.util.HashMap;
import java.util.Map;

/**
 * 非对称加密RSA
 *
 * @author jiangcs
 * @since 2022/08/25
 */
public class RsaUtil {
    private static final String transformation = "RSA";

    /**
     * 加密
     *
     * @param publicKey 公钥
     * @param plaintext 明文
     * @return 加密字符串
     * @throws Exception 加密异常
     */
    public static String encrypt(String publicKey, String plaintext) throws Exception {
        RSAPublicKey rsaPublicKey = (RSAPublicKey) KeyFactory.getInstance(Algorithm.RSA.code())
                .generatePublic(new X509EncodedKeySpec(publicKey.getBytes(StandardCharsets.UTF_8)));

        Cipher cipher = Cipher.getInstance(transformation);
        cipher.init(Cipher.ENCRYPT_MODE, rsaPublicKey);
        byte[] bytes = cipher.doFinal(plaintext.getBytes(StandardCharsets.UTF_8));
        return new String(bytes, StandardCharsets.UTF_8);
    }

    /**
     * 解密
     *
     * @param privateKey 私钥
     * @param ciphertext 密文
     * @return 解密字符串
     * @throws Exception 解密异常
     */
    public static String decrypt(String privateKey, String ciphertext) throws Exception {
        RSAPrivateKey rsaPrivateKey = (RSAPrivateKey) KeyFactory.getInstance(Algorithm.RSA.code())
                .generatePrivate(new PKCS8EncodedKeySpec(privateKey.getBytes(StandardCharsets.UTF_8)));

        Cipher cipher = Cipher.getInstance(transformation);
        cipher.init(Cipher.DECRYPT_MODE, rsaPrivateKey);
        byte[] bytes = cipher.doFinal(ciphertext.getBytes(StandardCharsets.UTF_8));
        return new String(bytes, StandardCharsets.UTF_8);
    }

    /**
     * 生成RSA公私钥
     *
     * @param keySize Key长度，1025 或 2048
     * @return 公钥，私钥
     * @throws Exception 生成密钥异常
     */
    public static Map<String, String> generateKeyPair(int keySize) throws Exception {
        KeyPairGenerator keyPairGenerator = KeyPairGenerator.getInstance(Algorithm.RSA.code());
        // 初始化密钥对生成器，密钥大小为1024 2048位
        keyPairGenerator.initialize(keySize, new SecureRandom());
        // 生成一个密钥对，保存在keyPair中
        KeyPair keyPair = keyPairGenerator.generateKeyPair();

        Map<String, String> result = new HashMap<>(2);
        // 得到公钥字符串
        result.put("publicKey", new String(Base64.getEncoder().encode(keyPair.getPublic().getEncoded())));
        // 得到私钥字符串
        result.put("privateKey", new String(Base64.getEncoder().encode(keyPair.getPrivate().getEncoded())));
        return result;
    }
}
